import tensorflow as tf
import numpy as np
from PIL import Image
import os


def get_all_captcha(dir):
    """
        得到所有的验证码文件路径
        :param dir:
        :return:
    """
    captcha_paths = []
    for root, dirs, files in os.walk(dir):
        for file in files:
            captcha_paths.append(os.path.join(root, file))
    return captcha_paths


def write_tfrecord(captcha_paths, tfrecord_path):
    """
        将图片的captcha转化为

    :param captcha_paths:
    :param tfrecord_path:
    :return:
    """
    with tf.Session() as session:
        with tf.python_io.TFRecordWriter(tfrecord_path) as write:
            for captcha_path in captcha_paths:
                try:
                    image_data = Image.open(captcha_path)

                    # 根据模型的结构resize
                    image_data = image_data.resize((224, 224))
                    # 灰度化
                    image_data = np.array(image_data.convert('L'))

                    # 将图片转化为bytes
                    image_data = image_data.tobytes()

                    liables = captcha_path.split("\\")[-1][:4]

                    # record记录向量
                    example = tf.train.Example(features=tf.train.Features(feature={

                        "image": tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_data])),
                        'label0': tf.train.Feature(int64_list=tf.train.Int64List(value=[int(liables[0])])),
                        'label1': tf.train.Feature(int64_list=tf.train.Int64List(value=[int(liables[1])])),
                        'label2': tf.train.Feature(int64_list=tf.train.Int64List(value=[int(liables[2])])),
                        'label3': tf.train.Feature(int64_list=tf.train.Int64List(value=[int(liables[3])])),
                    }))

                    write.write(example.SerializeToString())

                except ValueError as e:
                    print('Could not read:', captcha_path)
                    print('Error:', e)
                    print('Skip it\n')


if __name__ == '__main__':
    captcha_paths = get_all_captcha('F:\Idea workspace\\tensorflow-learning\data\sample\captcha')

    train_paths = captcha_paths[:int(len(captcha_paths) * 0.8)]
    test_paths = captcha_paths[int(len(captcha_paths) * 0.8):]
    write_tfrecord(train_paths, 'F:\Idea workspace\\tensorflow-learning\data\sample\\tfrecord\\captcha-train.tfrecord')
    write_tfrecord(test_paths, 'F:\Idea workspace\\tensorflow-learning\data\sample\\tfrecord\\captcha-test.tfrecord')
